from langchain_community.retrievers import SVMRetriever
from langchain_community.retrievers import TFIDFRetriever
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from tool import get_azure_endpoint,get_api_version,get_api_key
if __name__ == '__main__':
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    # 加载PDF
    loader_chinese = PyPDFLoader("./docs/matplotlib/第一回：Matplotlib初相识.pdf")
    pages_chinese = loader_chinese.load()
    all_page_text_chinese = [p.page_content for p in pages_chinese]
    joined_page_text_chinese = " ".join(all_page_text_chinese)
    # 分割文本
    text_splitter_chinese = RecursiveCharacterTextSplitter(chunk_size = 1500,chunk_overlap = 150)
    splits_chinese = text_splitter_chinese.split_text(joined_page_text_chinese)
    # 检索
    svm_retriever = SVMRetriever.from_texts(splits_chinese, embedding)
    tfidf_retriever = TFIDFRetriever.from_texts(splits_chinese)

    question_chinese = "这门课的主要主题是什么?"
    docs_svm_chinese = svm_retriever.get_relevant_documents(question_chinese)
    print(docs_svm_chinese[0])

    docs_tfidf_chinese = tfidf_retriever.get_relevant_documents(question_chinese)
    print(docs_tfidf_chinese[0])

